'''
Created on Aug 13, 2009

@author: mkiyer
'''

import sys
from bx.intervals.intersection import Interval, IntervalTree
import collections
import numpy as np

class GenomicInterval(object):    
    def __init__(self, chrom, start, end):
        self.chrom = chrom
        self.start = start
        self.end = end

    def parse_line(self, line):
        fields = line.strip().split('\t')
        self.chrom = fields[0]
        self.start = int(fields[1])
        self.end = int(fields[2])

class BedInterval(GenomicInterval):
    def __init__(self, chrom, start, end, name='.', score=1000, strand='+', 
                 tx_start=0, tx_end=0, item_rgb=0, 
                 block_count=0, block_sizes=None, block_starts=None):
        super(BedInterval, self).__init__(start, end, chrom=chrom, strand=strand)
        self.name = name
        self.score = score
        self.tx_start = tx_start
        self.tx_end = tx_end
        self.item_rgb = item_rgb
        self.block_count = block_count
        self.block_sizes = block_sizes
        self.block_starts = block_starts
    
    def parse_line(self, line):
        fields = line.strip().split('\t')
        self.chrom = fields[0]
        self.start = int(fields[1])
        self.end = int(fields[2])
        self.name = fields[3]
        self.score = int(fields[4])
        self.strand = fields[5]
        self.tx_start = int(fields[6])
        self.tx_end = int(fields[7])
        self.item_rgb = fields[8]
        self.block_count = int(fields[9])
        self.block_sizes = map(int, fields[10].split(','))
        self.block_starts = map(int, fields[11].split(','))

def parse_raw_interval_file(fhd, category):
    for id, line in enumerate(fhd):
        fields = line.strip().split('\t')
        chrom = fields[0]
        start = int(fields[1])
        end = int(fields[2])
        yield Interval(start, end, value=('%s_%d' % (category, id), category), chrom=chrom)
        
def parse_bed_file(fhd, category):
    for line in fhd:
        fields = line.strip().split('\t')
        chrom = fields[0]
        start = int(fields[1])
        end = int(fields[2])
        name = fields[3]
        strand = fields[5]
        yield Interval(start, end, value=(name, category), chrom=chrom, strand=strand)

def parse_kgtxinfo(fhd):
    kgtxinfo = {}
    for line in fhd:
        fields = line.strip().split('\t')
        kgtxinfo[fields[0]] = fields[1]
    return kgtxinfo

def parse_gene_pred(fhd, category):
    '''
    name varchar(255) NOT NULL default '',
    chrom varchar(255) NOT NULL default '',
    strand char(1) NOT NULL default '',
    txStart int(10) unsigned NOT NULL default '0',
    txEnd int(10) unsigned NOT NULL default '0',
    cdsStart int(10) unsigned NOT NULL default '0',
    cdsEnd int(10) unsigned NOT NULL default '0',
    exonCount int(10) unsigned NOT NULL default '0',
    exonStarts longblob NOT NULL,
    exonEnds longblob NOT NULL,
    proteinID varchar(40) NOT NULL default '',
    alignID varchar(255) NOT NULL default '',
    '''
    for line in fhd:
        fields = line.strip().split('\t')
        name = fields[0]
        chrom = fields[1]
        strand = fields[2]
        txstart = int(fields[3])
        txend = int(fields[4])
        yield Interval(txstart, txend, value=(name, category), chrom=chrom, strand=strand)

def load_ucsc_genes(kgfhd, kgtxinfofhd):
    kgtxinfodict = parse_kgtxinfo(kgtxinfofhd)
    # knowngene file is genePred format
    for interval in parse_gene_pred(kgfhd, 'ucsc'):
        ucsc_category = kgtxinfodict[interval.value[0]]
        if ucsc_category == "coding":
            interval.value = (interval.value[0], 'ucsc_coding')
        else:
            interval.value = (interval.value[0], 'ucsc_noncoding')
        #interval.value = (interval.value[0], 'ucsc_%s' % kgtxinfodict[interval.value[0]])
        yield interval
    del kgtxinfodict

def plot_cov_distributions(categories, covintervals, outfile):
    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm

    category_names = ['novel', 'c', 'nc', 'c+nc', 'l', 'c+l', 'nc+l', 'c+nc+l']
    scatters = []

    plt.figure(2)    
    for category, intervals in zip(categories, covintervals):
        # the data
        xdata = [x.cov_c for x in intervals]
        ydata = [x.cov_t for x in intervals]
        s = 20*np.log(np.array([x.ratio for x in intervals]))        
        #c = [x.category * cm.jet.N / len(category_names)] * len(intervals)
        c = x.category * cm.jet.N / len(category_names)
          
        plt.scatter(xdata, ydata, s=30, c=cm.hot(c), marker='o', cmap=None, norm=None,
                    vmin=None, vmax=None, alpha=1.0, linewidths=None,
                    verts=None, label=category_names[category])
    plt.legend()
    plt.xlabel('PrEC RNA-seq coverage (RPKM)')
    plt.ylabel('VCaP RNA-seq coverage (RPKM)')
    plt.savefig(outfile)
    
#    xdata = [x.cov_c for x in covintervals]
#    ydata = [x.cov_t for x in covintervals]
#    s = 20*np.log(np.array([x.ratio for x in covintervals]))
#    c = [x.category for x in covintervals]
#    
#    plt.figure(2)
#    plt.scatter(xdata, ydata, s=30, c=c, marker='o', cmap=None, norm=None,
#                vmin=None, vmax=None, alpha=0.75, linewidths=None,
#                verts=None)
#    plt.xlabel('PrEC RNA-seq coverage (RPKM)')
#    plt.ylabel('VCaP RNA-seq coverage (RPKM)')
#    plt.savefig(outfile)

def make_hist(categories):    
    # histogram
    #range=None, normed=False, weights=None)
    import matplotlib.pyplot as plt
    import pylab
    import numpy as np
    # the histogram of the data
    #n, bins, patches = plt.hist(exoncov, bins=50, log=True, facecolor='r', alpha=0.5)
    #n2, bins2, patches2 = plt.hist(covarray, bins, log=True, facecolor='g', alpha=0.5)
    #, normed=1, facecolor='g', alpha=0.75)
    #plt.legend([patches[0], patches2[0]], ['lander', 'all'])
    #n, bins, patches = plt.hist([categories], bins=16, range=(0, 50), log=True)    
    category_names = ['novel', 'c', 'nc', 'c+nc', 'l', 'c+l', 'nc+l', 'c+nc+l']
    nbins = 8
    bins = np.arange(nbins)
    n, bins, patches = plt.hist(categories, bins=nbins, range=(0,8), log=True)
    pos = np.arange(nbins)+0.5    #Center bars on the axis ticks
    pylab.xticks(pos, category_names)    
    plt.xlabel('category')
    plt.ylabel('# of intervals')
    plt.title('histogram of interval classes')
    plt.grid(True)
    #plt.axis([40, 160, 0, 0.03])
    plt.savefig('category.png')

def plot_hist(categories, category_names, outfile):
    import matplotlib.pyplot as plt
    import numpy as np
    cats = collections.defaultdict(lambda: 0)
    for cat in categories:
        cats[cat] += 1
    labels = sorted(cats.keys())
    bars = [cats[x] for x in labels]
    labels = [category_names[x] for x in labels]
    fig = plt.figure(figsize=(60, 15))
    plot_bar(labels, bars, outfile=outfile, fig=fig)
    return

def autolabel(ax, rects):
    # attach some text labels
    for rect in rects:
        height = rect.get_height()
        ax.text(rect.get_x()+rect.get_width()/2., 1.05*height, '%d'%int(height),
                ha='center', va='bottom')

def plot_bar(labels, bars, outfile, fig=None):
    import numpy as np
    import matplotlib.pyplot as plt
    width=0.8
    if fig == None:
        fig = plt.figure()
    ax = fig.add_subplot(111)
    ind = np.arange(len(labels))    
    rects1 = ax.bar(ind, bars, width)
    ax.set_xticks(ind)
    #ax.set_xticklabels(labels, rotation=45, fontsize=8)
    ax.set_xticklabels(labels)
    fig.autofmt_xdate()
    autolabel(ax, rects1)
    plt.savefig(outfile)

def plot_heatmap(heatmap, outfile, xticks=None):
    import matplotlib.pyplot as plt
    plt.figure(figsize=(25, 200))
    plt.imshow(heatmap[:,:], aspect='auto', interpolation='nearest')
    if xticks:
        plt.xticks(np.arange(heatmap.shape[1]), xticks)
    plt.grid(False)
    plt.savefig(outfile)


if __name__ == '__main__':
    # create the interval tree
    trees = collections.defaultdict(lambda: IntervalTree())
    allintervals = collections.defaultdict(lambda: [])
    
    # load all the different categories of intervals to compare
    #
    # UCSC genes
    #
    kgtxinfofhd = open('./externaldata/kgTxInfo.txt')
    kgfhd = open('./externaldata/knownGene.txt')    
    for interval in load_ucsc_genes(kgfhd, kgtxinfofhd):
        trees[interval.chrom].insert_interval(interval)
        allintervals[interval.value[1]].append(interval)
    #
    # Lander lincrnas
    #
    lander_fhd = open('./externaldata/lincrna_regions.txt')    
    for interval in parse_raw_interval_file(lander_fhd, "lincrna"):
        trees[interval.chrom].insert_interval(interval)
        allintervals[interval.value[1]].append(interval)
    #
    # PINS and TINS
    #
    pins_fhd = open('./externaldata/PINS_hg18.bed')
    tins_fhd = open('./externaldata/TINS_hg18.bed')
    for interval in parse_bed_file(pins_fhd, "PINS"):
        trees[interval.chrom].insert_interval(interval)
        allintervals[interval.value[1]].append(interval)
    for interval in parse_bed_file(tins_fhd, "TINS"):
        trees[interval.chrom].insert_interval(interval)
        allintervals[interval.value[1]].append(interval)
    #
    # RNAdb
    # 
    rnadb_antisense_fhd = open('./externaldata/hg18Homo_sapiensAntisense_ncRNA_Pipeline.bed')
    rnadb_mirna_fhd = open('./externaldata/hg18Homo_sapiensmiRNA.bed')
#    for interval in parse_bed_file(rnadb_antisense_fhd, "rnadb_antisense"):
    for interval in parse_bed_file(rnadb_antisense_fhd, "rnadb"):
        trees[interval.chrom].insert_interval(interval)
        allintervals[interval.value[1]].append(interval)
    for interval in parse_bed_file(rnadb_mirna_fhd, "rnadb"):
        trees[interval.chrom].insert_interval(interval)
        allintervals[interval.value[1]].append(interval)
    #
    # H3K36/H4 predictions
    #
    mydata = open(sys.argv[1])
    for interval in parse_raw_interval_file(mydata, "h3k4_h3k36"):
        trees[interval.chrom].insert_interval(interval)
        allintervals[interval.value[1]].append(interval)

    # primary concept category
    # TODO: set by command line
    primary_concept = "h3k4_h3k36"
    primary_intervals = allintervals[primary_concept]

    # determine all the different categories
    concepts = set(allintervals.keys())
    concepts.remove(primary_concept)
    concepts = sorted(list(concepts))
    concepts.insert(0, primary_concept)
    concept_lookup = {}
    index_lookup = {}
    concept_counts = {}
    for i, concept in enumerate(concepts):
        concept_lookup[concept] = i
        index_lookup[i] = concept
        concept_counts[concept] = 0

    categories = []
    category_names = []
    category_intervals = collections.defaultdict(lambda: [])
    for cat in xrange(0, 2**len(concepts)):
        thisname = []
        for n in xrange(len(concepts)):
            if (cat & (1 << n)) != 0:
                thisname.append(index_lookup[n])
        if len(thisname) == 0:
            thisname = ['none']
        category_names.append('+'.join(thisname))
    
    # create a matrix (heatmap) to store results
    heatmap = np.zeros((len(primary_intervals), len(concepts)), dtype=np.float)

    # user-defined set of desired intervals
    desired_concepts = {'ucsc_coding': False}

    for i, interval in enumerate(primary_intervals):
        if interval.chrom not in trees:
            continue
        
        # find all the overlapping intervals
        hits = trees[interval.chrom].find(interval.start, interval.end)
        hit_concepts = set([])   
        for hit in hits:
            hit_concept = hit.value[1]
            hit_concepts.add(hit_concept)
            heatmap[i, concept_lookup[hit_concept]] = 1

        # check if this interval the contraints of desired intervals
        desired_interval = True
        for concept_name, concept_value in desired_concepts.iteritems():
            if (concept_name in hit_concepts) != concept_value:
                desired_interval = False
                break
        if desired_interval:                
            sys.stdout.write('%s\t%d\t%d\t%s\n' % (interval.chrom,
                                                   interval.start,
                                                   interval.end,
                                                   ','.join(hit_concepts)))
        # count concept matches
        for concept in hit_concepts:
            concept_counts[concept] += 1
        # categorize by combinations of concepts
        category = 0
        for hit in hits:
            if hit.value[1] != primary_concept:
                category |= 1 << concept_lookup[hit.value[1]]
        categories.append(category)
    
    plot_hist(categories, category_names, "category.png")
    plot_bar(concepts, [concept_counts[x] for x in concepts], "count.png")
    plot_heatmap(heatmap, "heatmap.png", xticks=concepts)
    sys.exit(0)

    # some convenience functions to improve readability
    is_kg = lambda interval: interval.value.startswith('uc')
    is_kg_coding = lambda interval: kgtxinfo[interval.value] == 'coding'
    is_kg_noncoding = lambda interval: kgtxinfo[interval.value] != 'coding'
    is_lincrna = lambda interval: interval.value.startswith('lincrna')
    
    bin_by_category = collections.defaultdict(lambda: [])
    categories = []
    
    for interval in intervals:        
        if interval.chrom not in trees:
            continue        
        hits = trees[interval.chrom].find(interval.start, interval.end)
        ucsc = 0
        coding = 0
        noncoding = 0
        lincrna = 0

        for hit in hits:
            if is_kg(hit):
                ucsc = True
                coding |= is_kg_coding(hit)
                noncoding |= is_kg_noncoding(hit)            
            lincrna |= is_lincrna(hit)
        
        category = (lincrna << 2) | (noncoding << 1) | coding       
        categories.append(category)        
        interval.category = category
        bin_by_category[category].append(interval)
        
        sys.stdout.write("%s\t%d\t%d\t%d\t%d\t%s\n" %
                         (interval, int(ucsc), int(coding), int(noncoding),
                          int(lincrna), category))

    make_hist(categories)
    plot_cov_distributions([0,2,4,7], 
                           [bin_by_category[0], bin_by_category[2], bin_by_category[4], bin_by_category[7]], 
                           'novel_rnaseq_cov.png')


